package linker

import (
	"debug/elf"
	"math"
	"risc-v/ld/pkg/utils"
)

type OutPutProgramHeader struct {
	Chunk
	ProgramHeaders []ProgramHeader
}

func NewOutPutProgramHeader() *OutPutProgramHeader {
	o := &OutPutProgramHeader{Chunk: NewChunk()}
	o.SectionHeader.Flags = uint64(elf.SHF_ALLOC)
	o.SectionHeader.AddrAlign = 8

	return o
}

func toProgramHeaderFlags(chunk Chunker) uint32 {
	ret := uint32(elf.PF_R)

	write := chunk.GetSectionHeader().Flags&uint64(elf.SHF_WRITE) != 0
	if write {
		ret |= uint32(elf.PF_W)
	}

	if chunk.GetSectionHeader().Flags&uint64(elf.SHF_EXECINSTR) != 0 {
		ret |= uint32(elf.PF_X)
	}

	return ret
}

func createProgramHeader(ctx *Context) []ProgramHeader {
	vec := make([]ProgramHeader, 0)

	define := func(typ, flags uint64, minAlign int64, chunk Chunker) {
		vec = append(vec, ProgramHeader{})
		phdr := &vec[len(vec)-1]
		phdr.Type = uint32(typ)
		phdr.Flags = uint32(flags)
		phdr.Align = uint64(math.Max(float64(minAlign), float64(chunk.GetSectionHeader().AddrAlign)))
		phdr.Offset = chunk.GetSectionHeader().Offset

		if chunk.GetSectionHeader().Type == uint32(elf.SHT_NOBITS) {
			phdr.FileSize = 0
		} else {
			phdr.FileSize = chunk.GetSectionHeader().Size
		}

		phdr.VAddr = chunk.GetSectionHeader().Addr
		phdr.PAddr = chunk.GetSectionHeader().Addr
		phdr.MemSize = chunk.GetSectionHeader().Size
	}

	push := func(chunk Chunker) {
		phdr := &vec[len(vec)-1]
		phdr.Align = uint64(math.Max(float64(phdr.Align), float64(chunk.GetSectionHeader().AddrAlign)))

		if chunk.GetSectionHeader().Type != uint32(elf.SHT_NOBITS) {
			phdr.FileSize = chunk.GetSectionHeader().Addr + chunk.GetSectionHeader().Size - phdr.VAddr
		}

		phdr.MemSize = chunk.GetSectionHeader().Addr + chunk.GetSectionHeader().Size - phdr.VAddr
	}

	define(uint64(elf.PT_PHDR), uint64(elf.PF_R), 8, ctx.ProgramHeader)

	isTls := func(chunk Chunker) bool {
		return chunk.GetSectionHeader().Flags&uint64(elf.SHF_TLS) != 0
	}

	isBss := func(chunk Chunker) bool {
		return chunk.GetSectionHeader().Type == uint32(elf.SHT_NOBITS) && !isTls(chunk)
	}

	isNote := func(chunk Chunker) bool {
		shdr := chunk.GetSectionHeader()
		return shdr.Type == uint32(elf.SHT_NOTE) && shdr.Flags&uint64(elf.SHF_ALLOC) != 0
	}

	end := len(ctx.Chunks)
	for i := 0; i < end; {
		first := ctx.Chunks[i]
		i++
		if !isNote(first) {
			continue
		}

		flags := toProgramHeaderFlags(first)
		aligment := first.GetSectionHeader().AddrAlign

		define(uint64(elf.PT_NOTE), uint64(flags), int64(aligment), first)

		for i < end && isNote(ctx.Chunks[i]) && toProgramHeaderFlags(ctx.Chunks[i]) == flags {
			push(ctx.Chunks[i])
			i++
		}
	}

	{
		chunks := make([]Chunker, 0)

		chunks = append(chunks, ctx.Chunks...)

		chunks = utils.RemoveIf(chunks, func(chunk Chunker) bool {
			return isTbss(chunk)
		})

		end := len(chunks)
		for i := 0; i < end; {
			first := chunks[i]
			i++

			if first.GetSectionHeader().Flags&uint64(elf.SHF_ALLOC) == 0 {
				break
			}

			flags := toProgramHeaderFlags(first)
			define(uint64(elf.PT_LOAD), uint64(flags), PageSize, first)

			if !isBss(first) {
				for i < end && !isBss(chunks[i]) && toProgramHeaderFlags(chunks[i]) == flags {
					push(chunks[i])
					i++
				}
			}

			for i < end && isBss(chunks[i]) && toProgramHeaderFlags(chunks[i]) == flags {
				push(chunks[i])
				i++
			}
		}
	}

	for i := 0; i < len(ctx.Chunks); i++ {
		if !isTls(ctx.Chunks[i]) {
			continue
		}

		// PT_TLS 类型的 Segment
		define(uint64(elf.PT_TLS), uint64(toProgramHeaderFlags(ctx.Chunks[i])), 1, ctx.Chunks[i])

		i++

		for i < len(ctx.Chunks) && isTls(ctx.Chunks[i]) {
			push(ctx.Chunks[i])
			i++
		}

		// RISC-V 规定 TpAddr 的值为 PT_TLS 类型 segment 的虚拟地址
		phdr := vec[len(vec)-1]
		ctx.TpAddr = phdr.VAddr
	}

	return vec
}

func (o *OutPutProgramHeader) UpdateSectionHeader(ctx *Context) {
	o.ProgramHeaders = createProgramHeader(ctx)
	o.SectionHeader.Size = uint64(len(o.ProgramHeaders)) * uint64(ProgramHeaderSize)
}

func (o *OutPutProgramHeader) CopyBuf(ctx *Context) {
	utils.Write(ctx.Buf[o.SectionHeader.Offset:], o.ProgramHeaders)
}
